package org.tech4j.schedule.powerjob.slice;

import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.tech4j.common.utils.JsonUtils;
import org.tech4j.schedule.powerjob.constant.JobConstants;
import tech.powerjob.worker.core.ha.ProcessorTrackerStatusHolder;
import tech.powerjob.worker.core.processor.ProcessResult;
import tech.powerjob.worker.core.processor.TaskContext;
import tech.powerjob.worker.core.processor.TaskResult;
import tech.powerjob.worker.core.processor.sdk.MapReduceProcessor;
import tech.powerjob.worker.core.tracker.task.TaskTracker;
import tech.powerjob.worker.core.tracker.task.TaskTrackerPool;

import java.lang.reflect.Field;
import java.util.List;
import java.util.Map;
import java.util.stream.IntStream;

/**
 * @program: tech4j
 * @description:
 * @author: YuKai Fan
 * @create: 2025/3/10 09:11
 **/
public abstract class SliceProcessor implements MapReduceProcessor {
    private final static Logger log = LoggerFactory.getLogger(SliceProcessor.class);

    private static final String LOG_PREFIX = SliceProcessor.class.getName();

    @Override
    public ProcessResult process(TaskContext taskContext) throws Exception {
        if (isRootTask()) {
            int totalWorker = this.getTotalWorker();
            if (0 == totalWorker) {
                log.error(JobConstants.NO_WORKER_ERROR);
                return new ProcessResult(false, JobConstants.NO_WORKER_ERROR);
            }
            List<SliceTask> sliceTaskList = IntStream.range(0, totalWorker).mapToObj(index -> new SliceTask(index, totalWorker)).toList();
            try {
                this.map(sliceTaskList, JobConstants.PROCESS_SUB_TASK_NAME);
                log.info("{}, distribute subtask success, count is : {}", LOG_PREFIX, sliceTaskList.size());
            } catch (Exception e) {
                sliceTaskList.clear();
            }
            return new ProcessResult(true);
        }
        SliceTask subTask = (SliceTask) taskContext.getSubTask();
        subTask.setContext(taskContext);
        return this.execute(subTask);
    }

    @Override
    public ProcessResult reduce(TaskContext taskContext, List<TaskResult> list) {
        return this.after(taskContext, list);
    }

    public ProcessResult after(TaskContext context, List<TaskResult> results) {
        Map<String, String> errorMap = Maps.newHashMap();
        for (TaskResult result : results) {
            if (!result.isSuccess()) {
                errorMap.put(result.getTaskId(), result.getResult());
            }
        }
        return errorMap.isEmpty() ? new ProcessResult(true)
                : new ProcessResult(false, JsonUtils.toJsonString(errorMap));
    }

    /**
     * 分片任务处理器
     *
     * @param subTask
     * @return
     * @throws Exception
     */
    protected abstract ProcessResult execute(SliceTask subTask) throws Exception;

    private int getTotalWorker() throws NoSuchFieldException, IllegalAccessException {
        Field id2TaskTracker = TaskTrackerPool.class.getDeclaredField(JobConstants.TASK_TRACKER);
        id2TaskTracker.setAccessible(Boolean.TRUE);
        Map<Long, TaskTracker> taskTrackerMap = (Map<Long, TaskTracker>) id2TaskTracker.get(JobConstants.NULL);

        if (CollectionUtils.isEmpty(taskTrackerMap)) {
            return 0;
        }
        log.info("{}, trackerMap is : {}", LOG_PREFIX, JsonUtils.toJsonString(taskTrackerMap));
        TaskTracker taskTracker = taskTrackerMap.values().iterator().next();
        /*
        这里获取TaskTracker其实是子类的CommonTaskTracker
        需要反射获取父类的TaskTracker的状态管理器
         */
        Class<?> superclass = taskTracker.getClass().getSuperclass();
        Field holder = superclass.getDeclaredField(JobConstants.STATUS_HOLDER);
        holder.setAccessible(true);
        // 非private属性，需要有当前实例来获取属性
        ProcessorTrackerStatusHolder statusHolder = (ProcessorTrackerStatusHolder) holder.get(taskTracker);
        int totalWorkers = statusHolder.getAllProcessorTrackers().size();
        log.info("{}, current worker size is : {}", LOG_PREFIX, totalWorkers);
        return totalWorkers;
    }
}
